#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
回测模块
用于测试交易策略的表现
"""

import pandas as pd
import numpy as np


def backtest_strategy(signals, initial_capital=10000.0):
    """
    对交易策略进行回测
    
    参数:
    signals: 包含交易信号的数据
    initial_capital: 初始资金
    
    返回:
    dict: 回测结果
    """
    # 创建一个 portfolio DataFrame 来跟踪投资组合
    portfolio = signals.copy()
    
    # 计算持仓数量：假设每次交易都使用全部资金购买/卖出股票
    # 这里简化处理，每次信号出现时，如果有资金就买入，如果有持仓就卖出
    portfolio['holdings'] = np.floor(initial_capital / signals['close'])
    
    # 计算总资产价值
    portfolio['total_asset'] = initial_capital
    
    # 简化的回测逻辑：
    # 1. 当出现买入信号时，假设我们用全部资金买入股票
    # 2. 当出现卖出信号时，假设我们卖出所有股票
    # 3. 计算最终资产价值
    
    current_position = 0  # 当前持仓数量
    cash = initial_capital  # 当前现金
    
    for i, (index, row) in enumerate(portfolio.iterrows()):
        if row['positions'] == 1.0:  # 买入信号
            # 计算可以买入的股票数量
            shares_to_buy = int(cash / row['close'])
            # 更新现金和持仓
            cash -= shares_to_buy * row['close']
            current_position += shares_to_buy
            
        elif row['positions'] == -1.0:  # 卖出信号
            # 卖出所有持仓
            cash += current_position * row['close']
            current_position = 0
            
        # 计算当前总资产
        portfolio.loc[index, 'total_asset'] = cash + current_position * row['close']
    
    # 计算收益率
    portfolio['returns'] = portfolio['total_asset'].pct_change()
    
    # 计算最终资产和总收益率
    final_asset = portfolio['total_asset'].iloc[-1]
    total_return = (final_asset - initial_capital) / initial_capital * 100
    
    return {
        'initial_capital': initial_capital,
        'final_asset': final_asset,
        'total_return': total_return,
        'portfolio': portfolio
    }